#!/usr/bin/env python3
import subprocess
from fire import Fire
from queue import SimpleQueue
from glob import glob
from collections import defaultdict, Counter
from os.path import dirname, join
from datetime import date
import zd
import sys
from threading import Thread, Lock
utf8 = 'utf8'


def extract(egin, end, html):
  if not html:
    return
  if egin is None:
    start = 0
  else:
    start = html.find(egin)
    if start >= 0:
      start += len(egin) + 1
    else:
      return
  if end is not None:
    end = html.find(end, start)
  if end is None or end >= 0:
    return html[start:end].strip()


def zcat(infile):
  proc = subprocess.Popen(
    ['pigz', "-dc", infile], stdout=subprocess.PIPE, encoding=None
  )
  return proc.stdout


class Ex:
  def __init__(self, txt):
    self.txt = txt

  def __call__(self, key):
    return extract(b'"' + key + b'":', b'"', self.txt)


def day_count_by_gz(gzfile, exist, id2phone, callback):
  day_id = defaultdict(lambda: Counter())
  for pos, txt in enumerate(zcat(gzfile)):
    ex = Ex(txt)
    id = ex(b'AndroidId')
    if id not in exist and id not in id2phone:
      phone = ex(b'userId')
      if phone:
        id2phone[id] = [phone,str(date.fromtimestamp(int(ex(b'firstSessionTimestamp').decode('ascii'))//1000)).encode(utf8)]
      else:
        continue
    day = int(ex(b'sessionTimestamp') or 0) // 86400000
    day_id[day][id] += 1
    #if pos > 1000000:
      #break
  callback(day_id)


@Fire
def main(dirpath):
  with open("android.txt", "rb") as android:
    exist = set(i.strip() for i in android)

  process_li = []

  day_id = defaultdict(lambda: Counter())

  lock = Lock()
  def callback(d):
    with lock:
      for day, c in d.items():
        day_id[day] += c

  todo = []
  id2phone = {}
  for i in glob(join(dirpath, "*.json.gz")):
    print(i)
    p = Thread(
      target=day_count_by_gz, args=(i, exist, id2phone, callback)
    )
    todo.append(p)
    p.start()

  for i in todo:
    i.join()

  day_li = tuple(sorted(day_id.keys()))

  outfile = dirpath + ".csv.zst"
  with zd.open(outfile, "w") as out:
    out.write(
      b"phone,AndroidId," + b",".join(
        str(date.fromtimestamp(i * 86400)).encode(utf8)
        for i in day_li
      ) + b"\n"
    )
    for id in list(exist)+list(id2phone.keys()):
      li = [day_id[day][id] for day in day_li]
      if any(li):
        out.write(
          b",".join(
            id2phone.get(id, [b'',b''])+[ id] +
            list(map(lambda i: str(i).encode(utf8) if i else b'', li))
          ) + b"\n"
        )
